#include "MeanShift.h"

#include <set>
#include <QMap>

template <int N>
MeanShift<N>::MeanShift( std::vector< Point > &pts )
{
	// the data
	data_points = pts;

	// construct the kd tree
	kd_tree = new KDTree;
	foreach (Point p, data_points)
		kd_tree->addPoint(p);
	kd_tree->build();

	setParameters();
}


template <int N>
void MeanShift<N>::setParameters()
{
	computeAABB();
	double space_radius = (bbmax - bbmin).norm();

	h = 0.15 * space_radius;
	e = 0.001 * space_radius;
	// a mode needs at least 1 percent support
	rho = data_points.size() / 100.0;

	grid_stride = 0.1 * space_radius;
}



template <int N>
double MeanShift<N>::kernelEpanechnikov( double u )
{
	return (abs(u) > 1.0)? 0 : 0.75 * (1.0 - u * u);
}


template <int N>
void MeanShift<N>::computeAABB()
{
	bbmin = Point( FLT_MAX);
	bbmax = Point(-FLT_MAX);	

	foreach(Point p, data_points)
	{
		bbmin.minimize(p);
		bbmax.maximize(p);
	}	
}

template <int N>
void MeanShift<N>::createGrid()
{
	computeAABB();

	Point diff = bbmax - bbmin;

	grid_radix = Grid_Coord(0);
	for (int i = 0; i < dim(); i++)
	{
		grid_radix[i] = (int) std::ceil(diff[i] / grid_stride);
	}
}

template <int N>
bool MeanShift<N>::moveToNext( Grid_Coord &gcoord )
{
	int carry = 1;
	for (int d = dim() - 1; d >= 0; d--)
	{
		// Apply carry
		gcoord[d] += carry;

		// New carry happens
		if (gcoord[d] == grid_radix[d])
		{
			carry = 1;
			gcoord[d] = 0;
		}
		// Otherwise stop
		else
		{
			carry = 0;
			break;
		}
	}

	// carry == 1 means overflow
	return (carry == 0);
}

template <int N>
Vector<double, N> MeanShift<N>::gridPoint( Grid_Coord gcoord )
{
	Point p;

	for (int i = 0; i < dim(); i++)
	{
		double step = gcoord[i] + 0.5;
		p[i] = bbmin[i] + step * grid_stride;
	}

	return p;
}


template <int N>
Vector<double, N> MeanShift<N>::meanOf( KDResults &kd_results )
{
	Point mean(0.0);
	double total_w(0.0);
	foreach(KDResultPair kdrp, kd_results)
	{
		int nid = kdrp.first;
		double w = kernelEpanechnikov(kdrp.second / h);

		mean += w * data_points[nid];
		total_w += w;
	}

	return mean / total_w;
}


template <int N>
double MeanShift<N>::densityAt( Point p )
{
	KDResults kd_results;
	kd_tree->ball_search(p, h, kd_results);

	double density(0.0);
	foreach(KDResultPair kdrp, kd_results)
	{
		density += kernelEpanechnikov(kdrp.second / h);
	}

	return density;
}



template <int N>
void MeanShift<N>::seekCandidateModes()
{
	// initial
	candidate_modes.clear();
	candidate_tags.resize(data_points.size(), std::vector<int>());

	// starting from regular grids
	Point new_mean, old_mean;
	KDResults kd_results;
	int tag;

	createGrid();
	Grid_Coord gcoord(0);
	do 
	{
		new_mean = gridPoint(gcoord);
		tag = candidate_modes.size();

		// shift mean until convergence
		do 
		{
			old_mean =  new_mean;

			kd_tree->ball_search(old_mean, h, kd_results);
			if (!kd_results.empty())
			{
				// update the new mean
				new_mean = meanOf(kd_results);

				// tag the data points
				foreach(KDResultPair kdrp, kd_results)
				{
					candidate_tags[kdrp.first].push_back(tag);
				}
			}
		} while ((new_mean - old_mean).norm() > e);

		// one candidate mode
		candidate_modes.push_back(old_mean);

	} while (moveToNext(gcoord));

	// density for candidates
	candidate_density.clear();
	foreach (Point cm, candidate_modes)
		candidate_density.push_back(densityAt(cm));
}



template <int N>
void MeanShift<N>::mergeCandidateModes()
{
	modes.clear();

	// merging close modes
	std::vector<bool> is_merged(candidate_modes.size(), false);
	candidate_to_mode.resize(candidate_modes.size(), -1);

	int m_id = 0;	// the next un-merged candidate id
	do 
	{
		// search for the nearby candidate modes
		is_merged[m_id] = true;

		std::vector<int> nearbyCandidateIDs;
		nearbyCandidateIDs.push_back(m_id);

		for (int cm_id = m_id + 1; cm_id < (int) candidate_modes.size(); cm_id++)
		{
			Point dv = candidate_modes[m_id] - candidate_modes[cm_id];

			if (dv.norm() < h)
			{
				nearbyCandidateIDs.push_back(cm_id);
				is_merged[cm_id] = true;
			}
		}


		// merge to the one with max density
		double max_density = -1;
		int selected_id;
		foreach (int i, nearbyCandidateIDs)
		{
			if (candidate_density[i] > max_density)
			{
				max_density = candidate_density[i];
				selected_id = i;
			}
		}

		// merge only when the max density is above the threshold
		if (max_density > rho)
		{
			// map each candidate to the new mode
			int to_id = modes.size();
			foreach(int from_id, nearbyCandidateIDs)
			{
				candidate_to_mode[from_id] = to_id;
			}

			// the new mode
			modes.push_back(candidate_modes[selected_id]);
		}

		// The next mode index
		for (m_id++; m_id < (int)candidate_modes.size(); m_id++)
		{
			if (!is_merged[m_id])
				break;
		}

	}while(m_id < (int)candidate_modes.size());

}


template <int N>
void MeanShift<N>::cluster()
{
	seekCandidateModes();

	mergeCandidateModes();

	clusters.clear();

	// select the tag for each data point via majority voting
	clusters.resize(modes.size(), std::vector<int>());
	for (int pi = 0; pi < (int)candidate_tags.size(); pi++)
	{
		std::vector<int> &tags = candidate_tags[pi];

		for (int i = 0; i < (int)tags.size(); i++)
		{
			tags[i] = candidate_to_mode[tags[i]];
		}

		// find the majority
		if (tags.empty()) continue;
		int mode_id = majorityOf(tags);
		if (mode_id != -1)	
			clusters[mode_id].push_back(pi);
		//otherwise: the point is not covered by any of the mode, should not happen for using all grid vertices
	}
}


template <int N>
std::vector< Vector<double, N> > MeanShift<N>::pointsInCluster( int cid )
{
	std::vector<Point> pts;
	if (cid >= 0 && cid < clusters.size())
	{
		foreach (int pid, clusters[cid])
			pts.push_back(data_points[pid]);
	}

	return pts;
}



template <int N>
QSet<int> MeanShift<N>::pointIDsNearMode( int cid, int scale /*=1*/ )
{
	QSet<int> ids;

	if (cid >= 0 && cid < (int)modes.size())
	{
		KDResults kd_results;
		kd_tree->ball_search(modes[cid], h * scale * scale, kd_results);	// squared distance

		foreach (KDResultPair kdrp, kd_results)
			ids.insert(kdrp.first);
	}

	return ids;
}



template <int N>
int MeanShift<N>::majorityOf( std::vector<int> &item_array )
{
	// unique items
	std::set<int> item_set;
	foreach(int item, item_array) item_set.insert(item);

	// initial counters
	QMap<int, int> item_counter;
	foreach(int item, item_set) item_counter[item] = 0;

	// count
	foreach(int item, item_array) 
		item_counter[item]++;

	// the majority
	int majority_item;
	int max_counts = 0;
	foreach(int item, item_counter.keys())
	{
		if(item_counter[item] > max_counts)
		{
			max_counts = item_counter[item];
			majority_item = item;
		}
	}

	return majority_item;
}

